# -*- coding: utf-8 -*-

"""
Calculate TG data difference and output new schema
"""

from __future__ import unicode_literals
from __future__ import absolute_import
from __future__ import division

from kgpipeline.job import KgJob, JobInputMisMatchError
from kgpipeline.jsonutil import entity_json_decode, entity_json_encode
from kgpipeline.sparkutil import to_text_dataframe, json_to_text_dataframe, text_dataframe_to_json
from kgpipeline.calc.rdd_diff import rdd_diff_by_id, handle_diff
from kgpipeline.calc.diff import key_diff
import kgpipeline.tg as tg
import logging

from typing import Dict, Set, AnyStr, List

import json

logger = logging.getLogger(__name__)


def diff_handler_tg(dict1, dict2, ignore):  # type: (Dict, Dict, Set[AnyStr]) -> List[Dict]
    """
    TG发版用的diff handler。
    相对于Fix格式：
    1. 处理实体时，需要保留_id, _t字段，并且添加diff类型字段_diff
    2. 处理链接时，需要保留_id, _t, _it, _ot, _i, _o字段，并且添加diff类型字段_diff
    """
    if dict1 is None:
        res = {"_id": dict2["_id"], "_t": dict2["_t"], "_diff": "delete"}
        if "_i" in dict2:
            res["_i"] = dict2["_i"]
            res["_it"] = dict2["_it"]
        if "_o" in dict2:
            res["_o"] = dict2["_o"]
            res["_ot"] = dict2["_ot"]
        return [res]
    elif dict2 is None:
        res = dict1
        res["_diff"] = "add"
        return [res]
    else:
        key_remove, key_add, key_update = key_diff(dict1, dict2, ignore)
        if key_remove or key_add or key_update:  # if there is any change
            res = {"_id": dict1["_id"], "_t": dict1["_t"], "_diff": "update"}
            for key in key_remove:
                res[key] = None
            for key in key_add:
                res[key] = dict1[key]
            for key in key_update:
                res[key] = dict1[key]
            if "_i" in dict2:
                res["_i"] = dict2["_i"]
                res["_it"] = dict2["_it"]
            if "_o" in dict2:
                res["_o"] = dict2["_o"]
                res["_ot"] = dict2["_ot"]
            return [res]
        else:
            return []


def calc_diff(entity_old, entity_new):
    ignored = {"dateModified"}
    rdd = rdd_diff_by_id(entity_new, entity_old, id_field="_id")
    return handle_diff(rdd, ignored, diff_handler_tg)


class KgTgDiffJob(KgJob):
    def process(self, inputs):
        if len(inputs) != 6:
            raise JobInputMisMatchError("needs exactly 6 input files for Tiger Graph Data Loading Job")

        tg_schema_old = text_dataframe_to_json(inputs[0][0])
        tg_schema_new = text_dataframe_to_json(inputs[3][0])

        entity_old = inputs[1][0].rdd.map(lambda x: entity_json_decode(x["value"]))
        link_old = inputs[2][0].rdd.map(lambda x: entity_json_decode(x["value"]))

        entity_new = inputs[4][0].rdd.map(lambda x: entity_json_decode(x["value"]))
        link_new = inputs[5][0].rdd.map(lambda x: entity_json_decode(x["value"]))

        tg_schema_new, tg_schema_diff = tg.diff_tg_schema(tg_schema_old, tg_schema_new)

        entity_diff = calc_diff(entity_old, entity_new)
        link_diff = calc_diff(link_old, link_new)

        return [
            json_to_text_dataframe(tg_schema_new),
            json_to_text_dataframe(tg_schema_diff),
            to_text_dataframe(entity_diff, entity_json_encode),
            to_text_dataframe(link_diff, entity_json_encode)
        ]


KgJob.register("tgdiff", KgTgDiffJob)
